#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
股价与情感指数的指数移动平均(EMA)分析工具
用于使用EMA平滑数据，分析股票价格与舆论情感之间的趋势关系
"""

import os
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import numpy as np
from datetime import datetime
import seaborn as sns
import matplotlib as mpl
from matplotlib.font_manager import FontProperties

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False
sns.set(font='SimHei')

font_path = './font/SimHei.ttf'
font_prop = FontProperties(fname=font_path)
plt.rcParams['font.family'] = font_prop.get_name()
mpl.font_manager.fontManager.addfont(font_path)

class StockSentimentEMAAnalyzer:
    """股价与情感指数指数移动平均分析类"""
    
    def __init__(self, stock_data_path, sentiment_data_path, decay_factor=0.8):
        """
        初始化分析器
        
        参数:
            stock_data_path: 股价数据CSV文件路径
            sentiment_data_path: 情感指数CSV文件路径
            decay_factor: 时间衰减因子，用于计算权重（0~1之间，越接近1衰减越慢）
        """
        self.stock_data_path = stock_data_path
        self.sentiment_data_path = sentiment_data_path
        self.stock_data = None
        self.sentiment_data = None
        self.merged_data = None
        self.decay_factor = decay_factor
        
        # 加载数据
        self._load_data()
        
    def _load_data(self):
        """加载股价和情感指数数据"""
        # 检查文件是否存在
        if not os.path.exists(self.stock_data_path):
            raise FileNotFoundError(f"股价数据文件不存在: {self.stock_data_path}")
        
        if not os.path.exists(self.sentiment_data_path):
            raise FileNotFoundError(f"情感指数数据文件不存在: {self.sentiment_data_path}")
        
        # 加载股价数据
        self.stock_data = pd.read_csv(self.stock_data_path)
        print(f"成功加载股价数据，共 {len(self.stock_data)} 条记录")
        
        # 加载情感指数数据
        self.sentiment_data = pd.read_csv(self.sentiment_data_path)
        print(f"成功加载情感指数数据，共 {len(self.sentiment_data)} 条记录")
        
        # 确保日期列格式正确
        self._format_date_columns()
        
        # 合并数据
        self._merge_data(self.decay_factor)
        
    def _format_date_columns(self):
        """处理日期列格式"""
        # 对股价数据处理日期列
        if "日期" in self.stock_data.columns:
            self.stock_data.rename(columns={"日期": "date"}, inplace=True)
        
        # 确保日期列格式统一
        try:
            self.stock_data['date'] = pd.to_datetime(self.stock_data['date'])
            self.sentiment_data['date'] = pd.to_datetime(self.sentiment_data['date'])
        except Exception as e:
            print(f"日期格式转换错误: {str(e)}")
            # 如果转换失败，尝试另一种方式
            try:
                self.stock_data['date'] = pd.to_datetime(self.stock_data['date'], format='%Y-%m-%d')
                self.sentiment_data['date'] = pd.to_datetime(self.sentiment_data['date'], format='%Y-%m-%d')
            except Exception as e:
                print(f"第二次尝试日期格式转换错误: {str(e)}")
    
    def _apply_time_decay_weights(self, group, decay_factor=0.8):
        """
        对同组内的情感指数应用时间衰减权重
        
        参数:
            group: 同一交易日的情感数据组
            decay_factor: 衰减因子，数值越小衰减越快
            
        返回:
            加权平均后的情感指数
        """
        if len(group) <= 1:
            return group['sentiment_index'].mean()
        
        # 将组内数据按时间排序（最近的在后面）
        group = group.sort_values(by='date')
        
        # 计算每条数据距离交易日结束的天数差
        max_date = group['date'].max()
        group['days_diff'] = (max_date - group['date']).dt.total_seconds() / (24 * 3600)
        
        # 计算指数衰减权重：越接近交易日结束，权重越大
        group['weight'] = np.power(decay_factor, group['days_diff'])
        
        # 归一化权重
        group['weight'] = group['weight'] / group['weight'].sum()
        
        # 计算加权平均
        weighted_sentiment = (group['sentiment_index'] * group['weight']).sum()
        
        return weighted_sentiment
    
    def _aggregate_by_trading_day(self, decay_factor=0.8):
        """
        按交易日聚合情感数据，并应用时间衰减权重
        
        参数:
            decay_factor: 时间衰减因子
            
        返回:
            按交易日聚合后的情感数据
        """
        # 确保情感数据有日期列
        if 'date' not in self.sentiment_data.columns:
            raise ValueError("情感指数数据缺少date列")
        
        # 提取股价数据中的交易日
        trading_days = set(self.stock_data['date'].dt.date)
        
        # 为情感数据添加一个交易日列，映射到最近的交易日
        self.sentiment_data['trading_date'] = self.sentiment_data['date'].apply(
            lambda x: self._find_nearest_trading_day(x, trading_days)
        )
        
        # 按交易日分组，应用时间衰减权重
        aggregated_sentiment = self.sentiment_data.groupby('trading_date').apply(
            lambda x: self._apply_time_decay_weights(x, decay_factor)
        ).reset_index()
        
        # 重命名列
        aggregated_sentiment.columns = ['date', 'sentiment_index']
        
        # 转换回datetime格式
        aggregated_sentiment['date'] = pd.to_datetime(aggregated_sentiment['date'])
        
        return aggregated_sentiment
    
    def _find_nearest_trading_day(self, date, trading_days):
        """
        找到最近的交易日
        
        参数:
            date: 日期
            trading_days: 交易日集合
            
        返回:
            最近的交易日
        """
        date_obj = date.date()
        
        # 如果当前日期是交易日，直接返回
        if date_obj in trading_days:
            return date_obj
        
        # 查找之前的交易日（非交易时段发帖归属到之前的交易日）
        prev_days = [day for day in trading_days if day < date_obj]
        if prev_days:
            return max(prev_days)
        
        # 如果没有之前的交易日，则归属到之后最近的交易日
        next_days = [day for day in trading_days if day > date_obj]
        if next_days:
            return min(next_days)
        
        # 如果都没有，返回原日期
        return date_obj
    
    def _merge_data(self, decay_factor=0.8):
        """
        合并股价和情感指数数据，应用交易日聚合和时间衰减权重
        
        参数:
            decay_factor: 时间衰减因子
        """
        # 确保两个数据集都有日期列
        if 'date' not in self.stock_data.columns:
            raise ValueError("股价数据缺少date列")
        
        # 处理股价数据中的列名
        if "收盘" in self.stock_data.columns:
            self.stock_data.rename(columns={"收盘": "close_price"}, inplace=True)
        elif "Close" in self.stock_data.columns:
            self.stock_data.rename(columns={"Close": "close_price"}, inplace=True)
        
        # 按交易日聚合情感数据，并应用时间衰减权重
        aggregated_sentiment = self._aggregate_by_trading_day(decay_factor)
        
        # 使用日期列合并数据
        self.merged_data = pd.merge(
            self.stock_data[['date', 'close_price']], 
            aggregated_sentiment, 
            on='date', 
            how='inner'
        )
        
        # 将日期设置为索引
        self.merged_data.set_index('date', inplace=True)
        
        # 按日期排序
        self.merged_data.sort_index(inplace=True)
        self.merged_data = self.merged_data[-125:]  # 取最近125天的数据
        
        # 添加归一化处理的数据列
        self._add_normalized_columns()
        
        # 计算指数移动平均
        self._calculate_ema()
        
        print(f"数据合并完成，共 {len(self.merged_data)} 条记录")
        print(f"应用了交易日聚合、时间衰减权重和指数移动平均")
        print(self.merged_data.head())
    
    def _add_normalized_columns(self):
        """添加归一化处理的数据列"""
        if self.merged_data is None or len(self.merged_data) == 0:
            return
        
        # 1. Min-Max归一化 (0-1区间)
        self.merged_data['close_price_norm_minmax'] = (self.merged_data['close_price'] - self.merged_data['close_price'].min()) / (self.merged_data['close_price'].max() - self.merged_data['close_price'].min())
        self.merged_data['sentiment_norm_minmax'] = (self.merged_data['sentiment_index'] - self.merged_data['sentiment_index'].min()) / (self.merged_data['sentiment_index'].max() - self.merged_data['sentiment_index'].min())
        
        # 2. Z-score标准化 (均值0，标准差1)
        self.merged_data['close_price_norm_zscore'] = (self.merged_data['close_price'] - self.merged_data['close_price'].mean()) / self.merged_data['close_price'].std()
        self.merged_data['sentiment_norm_zscore'] = (self.merged_data['sentiment_index'] - self.merged_data['sentiment_index'].mean()) / self.merged_data['sentiment_index'].std()
        
        # 3. 相对变化率
        self.merged_data['close_price_pct_change'] = self.merged_data['close_price'].pct_change() * 100
        self.merged_data['sentiment_pct_change'] = self.merged_data['sentiment_index'].pct_change() * 100
    
    def _calculate_ema(self):
        """计算不同时间窗口的指数移动平均"""
        # 定义不同的时间窗口
        windows = [3, 5, 10, 20]
        
        # 计算股价的指数移动平均
        for window in windows:
            self.merged_data[f'close_price_ema{window}'] = self.merged_data['close_price'].ewm(span=window, adjust=False).mean()
            self.merged_data[f'sentiment_ema{window}'] = self.merged_data['sentiment_index'].ewm(span=window, adjust=False).mean()
            
            # 对归一化数据也计算EMA
            self.merged_data[f'close_price_norm_zscore_ema{window}'] = self.merged_data['close_price_norm_zscore'].ewm(span=window, adjust=False).mean()
            self.merged_data[f'sentiment_norm_zscore_ema{window}'] = self.merged_data['sentiment_norm_zscore'].ewm(span=window, adjust=False).mean()
            
            # 对变化率计算EMA
            self.merged_data[f'close_price_pct_change_ema{window}'] = self.merged_data['close_price_pct_change'].ewm(span=window, adjust=False).mean()
            self.merged_data[f'sentiment_pct_change_ema{window}'] = self.merged_data['sentiment_pct_change'].ewm(span=window, adjust=False).mean()
    
    def plot_ema_trend(self, window=5, title=None, figsize=(14, 8), save_path=None, dpi=100):
        """
        绘制股价和情感指数的指数移动平均走势图
        
        参数:
            window: EMA时间窗口
            title: 图表标题
            figsize: 图表大小
            save_path: 保存路径，如果为None则不保存
            dpi: 图像分辨率
        """
        if self.merged_data is None or len(self.merged_data) == 0:
            print("没有可用的数据进行绘图")
            return
        
        # 设置默认标题
        if title is None:
            title = f"股价与情感指数EMA({window})走势图"
        
        # 创建图表
        fig, ax1 = plt.subplots(figsize=figsize)
        
        # 绘制股价曲线和EMA
        color = 'tab:red'
        ax1.set_xlabel('日期')
        ax1.set_ylabel('股价', color=color)
        ax1.plot(self.merged_data.index, self.merged_data['close_price'], color=color, linewidth=1, alpha=0.5, label='股价(原始)')
        ax1.plot(self.merged_data.index, self.merged_data[f'close_price_ema{window}'], color='darkred', linewidth=2, label=f'股价(EMA{window})')
        ax1.tick_params(axis='y', labelcolor=color)
        
        # 设置x轴日期格式
        ax1.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
        plt.xticks(rotation=45)
        
        # 创建第二个y轴
        ax2 = ax1.twinx()
        color = 'tab:blue'
        ax2.set_ylabel('情感指数', color=color)
        ax2.plot(self.merged_data.index, self.merged_data['sentiment_index'], color=color, linewidth=1, alpha=0.5, label='情感指数(原始)')
        ax2.plot(self.merged_data.index, self.merged_data[f'sentiment_ema{window}'], color='darkblue', linewidth=2, label=f'情感指数(EMA{window})')
        ax2.tick_params(axis='y', labelcolor=color)
        
        # 添加网格线
        ax1.grid(True, linestyle='--', alpha=0.6)
        
        # 添加图例
        lines1, labels1 = ax1.get_legend_handles_labels()
        lines2, labels2 = ax2.get_legend_handles_labels()
        ax1.legend(lines1 + lines2, labels1 + labels2, loc='upper left')
        
        # 计算EMA之间的相关系数
        corr = self.merged_data[f'close_price_ema{window}'].corr(self.merged_data[f'sentiment_ema{window}'])
        ax2.annotate(f"EMA相关系数: {corr:.4f}", xy=(0.02, 0.95), xycoords='axes fraction',
                   fontsize=10, bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="gray", alpha=0.8))
        
        # 设置标题
        plt.title(title, fontsize=15)
        
        # 调整布局
        plt.tight_layout()
        
        # 保存图表
        if save_path:
            plt.savefig(save_path, dpi=dpi, bbox_inches='tight')
            print(f"EMA走势图已保存至: {save_path}")
        
        # 显示图表
        plt.show()
        
        return corr
    
    def plot_multi_ema_comparison(self, figsize=(16, 12), save_path=None, dpi=100):
        """
        绘制多个时间窗口EMA的对比图
        
        参数:
            figsize: 图表大小
            save_path: 保存路径，如果为None则不保存
            dpi: 图像分辨率
        """
        windows = [3, 5, 10, 20]
        
        # 创建2x2的子图
        fig, axes = plt.subplots(2, 2, figsize=figsize)
        axes = axes.flatten()
        
        for i, window in enumerate(windows):
            ax = axes[i]
            
            # 绘制股价EMA
            color = 'tab:red'
            ax.set_xlabel('日期')
            ax.set_ylabel('股价EMA', color=color)
            ax.plot(self.merged_data.index, self.merged_data[f'close_price_ema{window}'], color=color, linewidth=2)
            ax.tick_params(axis='y', labelcolor=color)
            
            # 设置x轴日期格式
            ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
            plt.setp(ax.xaxis.get_majorticklabels(), rotation=45)
            
            # 创建第二个y轴
            ax2 = ax.twinx()
            color = 'tab:blue'
            ax2.set_ylabel('情感指数EMA', color=color)
            ax2.plot(self.merged_data.index, self.merged_data[f'sentiment_ema{window}'], color=color, linewidth=2)
            ax2.tick_params(axis='y', labelcolor=color)
            
            # 添加网格线
            ax.grid(True, linestyle='--', alpha=0.6)
            
            # 计算相关系数
            corr = self.merged_data[f'close_price_ema{window}'].corr(self.merged_data[f'sentiment_ema{window}'])
            ax.set_title(f'EMA{window} (相关系数: {corr:.4f})')
        
        plt.suptitle("不同时间窗口EMA对比分析", fontsize=16)
        plt.tight_layout()
        plt.subplots_adjust(top=0.92)
        
        # 保存图表
        if save_path:
            plt.savefig(save_path, dpi=dpi, bbox_inches='tight')
            print(f"多EMA对比图已保存至: {save_path}")
        
        plt.show()
    
    def plot_ema_normalized(self, window=5, normalization_type='zscore', title=None, figsize=(14, 8), save_path=None, dpi=100):
        """
        绘制归一化后的EMA走势图
        
        参数:
            window: EMA时间窗口
            normalization_type: 归一化类型，可选'zscore'或'pct_change'
            title: 图表标题
            figsize: 图表大小
            save_path: 保存路径，如果为None则不保存
            dpi: 图像分辨率
        """
        if self.merged_data is None or len(self.merged_data) == 0:
            print("没有可用的数据进行绘图")
            return
        
        # 选择对应的归一化数据列
        if normalization_type == 'zscore':
            price_col = f'close_price_norm_zscore_ema{window}'
            sentiment_col = f'sentiment_norm_zscore_ema{window}'
            type_desc = "Z-score标准化"
            ylabel = "标准化值 (z-score)"
        elif normalization_type == 'pct_change':
            price_col = f'close_price_pct_change_ema{window}'
            sentiment_col = f'sentiment_pct_change_ema{window}'
            type_desc = "变化率"
            ylabel = "变化率 (%)"
        else:
            print(f"不支持的归一化类型: {normalization_type}")
            return
        
        # 过滤掉NaN值
        plot_data = self.merged_data.dropna(subset=[price_col, sentiment_col])
        
        # 如果没有提供标题，生成默认标题
        if title is None:
            title = f"股价与情感指数EMA({window})对比图 ({type_desc})"
        
        # 创建图表
        plt.figure(figsize=figsize)
        
        # 绘制归一化后的曲线
        plt.plot(plot_data.index, plot_data[price_col], 'r-', label=f'股价EMA{window}', linewidth=2)
        plt.plot(plot_data.index, plot_data[sentiment_col], 'b-', label=f'情感指数EMA{window}', linewidth=2)
        
        # 设置x轴日期格式
        plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
        plt.xticks(rotation=45)
        
        # 添加网格线和图例
        plt.grid(True, linestyle='--', alpha=0.6)
        plt.legend(loc='best')
        
        # 设置标题和标签
        plt.title(title, fontsize=15)
        plt.xlabel('日期')
        plt.ylabel(ylabel)
        
        # 添加相关系数注释
        corr = plot_data[price_col].corr(plot_data[sentiment_col])
        plt.annotate(f"相关系数: {corr:.4f}", xy=(0.02, 0.95), xycoords='axes fraction',
                   fontsize=12, bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="gray", alpha=0.8))
        
        # 调整布局
        plt.tight_layout()
        
        # 保存图表
        if save_path:
            plt.savefig(save_path, dpi=dpi, bbox_inches='tight')
            print(f"归一化EMA走势图已保存至: {save_path}")
        
        # 显示图表
        plt.show()
        
        return corr
    
    def ema_lag_analysis(self, window=5, max_lag=5, normalization_type='zscore', figsize=(14, 6), save_path=None, dpi=100):
        """
        EMA滞后相关性分析
        
        参数:
            window: EMA时间窗口
            max_lag: 最大滞后天数
            normalization_type: 归一化类型
            figsize: 图表大小
            save_path: 保存路径，如果为None则不保存
            dpi: 图像分辨率
        """
        if self.merged_data is None or len(self.merged_data) <= max_lag:
            print("数据不足，无法进行滞后分析")
            return
        
        # 选择对应的归一化数据列
        if normalization_type == 'zscore':
            price_col = f'close_price_norm_zscore_ema{window}'
            sentiment_col = f'sentiment_norm_zscore_ema{window}'
            type_desc = "Z-score标准化"
        elif normalization_type == 'pct_change':
            price_col = f'close_price_pct_change_ema{window}'
            sentiment_col = f'sentiment_pct_change_ema{window}'
            type_desc = "变化率"
        else:
            print(f"不支持的归一化类型: {normalization_type}")
            return
        
        # 准备数据
        data = self.merged_data.dropna(subset=[price_col, sentiment_col])
        
        # 计算不同滞后期的相关系数
        correlations = []
        for i in range(max_lag + 1):
            if i == 0:
                # 当前相关性
                corr = data[price_col].corr(data[sentiment_col])
            else:
                # 滞后相关性 (股价滞后于情感)
                corr = data[price_col].corr(data[sentiment_col].shift(-i))
            correlations.append(corr)
        
        # 创建图表
        plt.figure(figsize=figsize)
        
        plt.bar(range(max_lag + 1), correlations, color='skyblue', alpha=0.7)
        plt.plot(range(max_lag + 1), correlations, 'ro-', linewidth=2)
        
        plt.xlabel('滞后天数')
        plt.ylabel('相关系数')
        plt.title(f'EMA({window})股价与情感指数的滞后相关性分析 ({type_desc})')
        plt.xticks(range(max_lag + 1))
        plt.axhline(y=0, color='k', linestyle='--', alpha=0.3)
        
        # 添加数值标签
        for i, v in enumerate(correlations):
            plt.text(i, v + 0.02 * (1 if v >= 0 else -1), f'{v:.3f}', ha='center')
        
        plt.grid(True, linestyle='--', alpha=0.6)
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=dpi, bbox_inches='tight')
            print(f"EMA滞后相关性分析图已保存至: {save_path}")
        
        plt.show()
        
        # 找出最大相关系数对应的滞后期
        max_corr_idx = np.argmax(np.abs(correlations))
        max_corr = correlations[max_corr_idx]
        print(f"EMA({window})最强相关性在滞后 {max_corr_idx} 天时出现，相关系数为 {max_corr:.4f}")
        
        return correlations
    
    def compare_raw_vs_ema(self, window=5, normalization_type='zscore', figsize=(16, 12), save_path=None, dpi=100):
        """
        比较原始数据与EMA处理后数据的效果
        
        参数:
            window: EMA时间窗口
            normalization_type: 归一化类型
            figsize: 图表大小
            save_path: 保存路径，如果为None则不保存
            dpi: 图像分辨率
        """
        if self.merged_data is None or len(self.merged_data) == 0:
            print("没有可用的数据进行对比")
            return
        
        # 选择对应的归一化数据列
        if normalization_type == 'zscore':
            raw_price_col = 'close_price_norm_zscore'
            raw_sentiment_col = 'sentiment_norm_zscore'
            ema_price_col = f'close_price_norm_zscore_ema{window}'
            ema_sentiment_col = f'sentiment_norm_zscore_ema{window}'
            type_desc = "Z-score标准化"
            ylabel = "标准化值 (z-score)"
        elif normalization_type == 'pct_change':
            raw_price_col = 'close_price_pct_change'
            raw_sentiment_col = 'sentiment_pct_change'
            ema_price_col = f'close_price_pct_change_ema{window}'
            ema_sentiment_col = f'sentiment_pct_change_ema{window}'
            type_desc = "变化率"
            ylabel = "变化率 (%)"
        else:
            print(f"不支持的归一化类型: {normalization_type}")
            return
        
        # 过滤掉NaN值
        data = self.merged_data.dropna()
        
        # 创建图表
        fig, axes = plt.subplots(2, 2, figsize=figsize)
        
        # 1. 原始数据走势图
        ax1 = axes[0, 0]
        
        ax1.plot(data.index, data[raw_price_col], 'r-', label='股价', linewidth=2)
        ax1.plot(data.index, data[raw_sentiment_col], 'b-', label='情感指数', linewidth=2)
        
        ax1.set_xlabel('日期')
        ax1.set_ylabel(ylabel)
        ax1.set_title('原始数据')
        ax1.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
        plt.setp(ax1.xaxis.get_majorticklabels(), rotation=45)
        ax1.grid(True, linestyle='--', alpha=0.6)
        ax1.legend(loc='best')
        
        raw_corr = data[raw_price_col].corr(data[raw_sentiment_col])
        ax1.annotate(f"相关系数: {raw_corr:.4f}", xy=(0.02, 0.95), xycoords='axes fraction',
                   fontsize=10, bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="gray", alpha=0.8))
        
        # 2. EMA处理后走势图
        ax2 = axes[0, 1]
        
        ax2.plot(data.index, data[ema_price_col], 'r-', label=f'股价EMA{window}', linewidth=2)
        ax2.plot(data.index, data[ema_sentiment_col], 'b-', label=f'情感指数EMA{window}', linewidth=2)
        
        ax2.set_xlabel('日期')
        ax2.set_ylabel(ylabel)
        ax2.set_title(f'EMA({window})处理后')
        ax2.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
        plt.setp(ax2.xaxis.get_majorticklabels(), rotation=45)
        ax2.grid(True, linestyle='--', alpha=0.6)
        ax2.legend(loc='best')
        
        ema_corr = data[ema_price_col].corr(data[ema_sentiment_col])
        ax2.annotate(f"相关系数: {ema_corr:.4f}", xy=(0.02, 0.95), xycoords='axes fraction',
                   fontsize=10, bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="gray", alpha=0.8))
        
        # 3. 原始数据滞后分析
        ax3 = axes[1, 0]
        max_lag = 5
        
        # 计算不同滞后期的相关系数
        raw_correlations = []
        for i in range(max_lag + 1):
            if i == 0:
                corr = data[raw_price_col].corr(data[raw_sentiment_col])
            else:
                corr = data[raw_price_col].corr(data[raw_sentiment_col].shift(-i))
            raw_correlations.append(corr)
        
        ax3.bar(range(max_lag + 1), raw_correlations, color='skyblue', alpha=0.7)
        ax3.plot(range(max_lag + 1), raw_correlations, 'ro-', linewidth=2)
        
        ax3.set_xlabel('滞后天数')
        ax3.set_ylabel('相关系数')
        ax3.set_title('原始数据滞后分析')
        ax3.set_xticks(range(max_lag + 1))
        ax3.axhline(y=0, color='k', linestyle='--', alpha=0.3)
        ax3.grid(True, linestyle='--', alpha=0.6)
        
        # 4. EMA数据滞后分析
        ax4 = axes[1, 1]
        
        # 计算EMA不同滞后期的相关系数
        ema_correlations = []
        for i in range(max_lag + 1):
            if i == 0:
                corr = data[ema_price_col].corr(data[ema_sentiment_col])
            else:
                corr = data[ema_price_col].corr(data[ema_sentiment_col].shift(-i))
            ema_correlations.append(corr)
        
        ax4.bar(range(max_lag + 1), ema_correlations, color='skyblue', alpha=0.7)
        ax4.plot(range(max_lag + 1), ema_correlations, 'ro-', linewidth=2)
        
        ax4.set_xlabel('滞后天数')
        ax4.set_ylabel('相关系数')
        ax4.set_title(f'EMA({window})数据滞后分析')
        ax4.set_xticks(range(max_lag + 1))
        ax4.axhline(y=0, color='k', linestyle='--', alpha=0.3)
        ax4.grid(True, linestyle='--', alpha=0.6)
        
        # 找出最大相关系数
        raw_max_idx = np.argmax(np.abs(raw_correlations))
        raw_max_corr = raw_correlations[raw_max_idx]
        
        ema_max_idx = np.argmax(np.abs(ema_correlations))
        ema_max_corr = ema_correlations[ema_max_idx]
        
        ax3.annotate(f"最大: 滞后{raw_max_idx}天, 相关系数={raw_max_corr:.3f}", 
                    xy=(0.05, 0.05), xycoords='axes fraction', fontsize=9,
                    bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="gray", alpha=0.8))
        
        ax4.annotate(f"最大: 滞后{ema_max_idx}天, 相关系数={ema_max_corr:.3f}", 
                    xy=(0.05, 0.05), xycoords='axes fraction', fontsize=9,
                    bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="gray", alpha=0.8))
        
        # 设置主标题
        plt.suptitle(f"原始数据 vs EMA({window})处理后数据对比分析 ({type_desc})", fontsize=16)
        plt.tight_layout()
        plt.subplots_adjust(top=0.93)
        
        # 保存图表
        if save_path:
            plt.savefig(save_path, dpi=dpi, bbox_inches='tight')
            print(f"原始与EMA对比分析图已保存至: {save_path}")
        
        plt.show()
        
        # 打印对比结果
        print(f"原始数据 - 当前相关系数: {raw_corr:.4f}, 最大相关系数: {raw_max_corr:.4f} (滞后 {raw_max_idx} 天)")
        print(f"EMA({window})处理后 - 当前相关系数: {ema_corr:.4f}, 最大相关系数: {ema_max_corr:.4f} (滞后 {ema_max_idx} 天)")
        
        # 计算改进效果
        corr_improvement = (ema_corr - raw_corr) / abs(raw_corr) * 100 if raw_corr != 0 else float('inf')
        max_corr_improvement = (ema_max_corr - raw_max_corr) / abs(raw_max_corr) * 100 if raw_max_corr != 0 else float('inf')
        
        print(f"相关系数改进: {corr_improvement:.2f}%, 最大相关系数改进: {max_corr_improvement:.2f}%")
        
        return {
            'raw': {
                'correlation': raw_corr,
                'lag_correlations': raw_correlations,
                'max_lag': raw_max_idx,
                'max_correlation': raw_max_corr
            },
            'ema': {
                'correlation': ema_corr,
                'lag_correlations': ema_correlations,
                'max_lag': ema_max_idx,
                'max_correlation': ema_max_corr
            },
            'improvement': {
                'correlation': corr_improvement,
                'max_correlation': max_corr_improvement
            }
        }

def main():
    """主函数"""
    # 文件路径
    stock_data_path = "../data/stock_data/hk01810_stock_data(3).csv"
    sentiment_data_path = "../data/sentiment_index2.csv"
    
    # 创建EMA分析器，使用时间衰减因子0.9
    analyzer = StockSentimentEMAAnalyzer(stock_data_path, sentiment_data_path, decay_factor=0.9)
    
    # 1. 绘制不同EMA窗口的走势图
    print("\n=== 不同EMA窗口的股价与情感指数走势图 ===")
    for window in [3, 5, 10, 20]:
        corr = analyzer.plot_ema_trend(
            window=window,
            title=f"小米集团(01810)股价与情感指数EMA({window})走势图",
            save_path=f"result_ema/stock_sentiment_ema{window}_trend.png"
        )
        print(f"EMA({window}) 相关系数: {corr:.4f}")
    
    # 2. 多EMA窗口对比
    print("\n=== 多EMA窗口对比分析 ===")
    analyzer.plot_multi_ema_comparison(save_path="result_ema/multi_ema_comparison.png")
    
    # 3. 使用归一化的EMA数据
    print("\n=== 归一化后的EMA走势分析 ===")
    for norm_type in ['zscore', 'pct_change']:
        for window in [5, 10]:
            corr = analyzer.plot_ema_normalized(
                window=window,
                normalization_type=norm_type,
                save_path=f"result_ema/ema{window}_{norm_type}_normalized.png"
            )
            print(f"EMA({window}) {norm_type}归一化 相关系数: {corr:.4f}")
    
    # 4. EMA滞后分析
    print("\n=== EMA滞后相关性分析 ===")
    for window in [5, 10]:
        analyzer.ema_lag_analysis(
            window=window,
            normalization_type='zscore',
            save_path=f"result_ema/ema{window}_lag_analysis.png"
        )
    
    # 5. 原始数据与EMA处理后的对比
    print("\n=== 原始数据与EMA处理后的对比分析 ===")
    for window in [5, 10]:
        analyzer.compare_raw_vs_ema(
            window=window,
            normalization_type='zscore',
            save_path=f"result_ema/raw_vs_ema{window}_comparison.png"
        )
    
    # 比较不同EMA窗口的效果
    print("\n=== 不同EMA窗口效果对比 ===")
    results = {}
    
    for window in [3, 5, 10, 20]:
        data = analyzer.merged_data.dropna()
        raw_corr = data['close_price_norm_zscore'].corr(data['sentiment_norm_zscore'])
        ema_corr = data[f'close_price_norm_zscore_ema{window}'].corr(data[f'sentiment_norm_zscore_ema{window}'])
        
        improvement = (ema_corr - raw_corr) / abs(raw_corr) * 100 if raw_corr != 0 else float('inf')
        
        results[window] = {
            'raw_corr': raw_corr,
            'ema_corr': ema_corr,
            'improvement': improvement
        }
        
        print(f"EMA({window}) - 原始相关性: {raw_corr:.4f}, EMA相关性: {ema_corr:.4f}, 改进: {improvement:.2f}%")
    
    # 找出最佳EMA窗口
    best_window = max(results.keys(), key=lambda w: results[w]['ema_corr'])
    print(f"\n最佳EMA窗口是 {best_window}，相关系数为 {results[best_window]['ema_corr']:.4f}，改进 {results[best_window]['improvement']:.2f}%")

if __name__ == "__main__":
    main() 